#!/usr/bin/env python3
import rospy
import cv2
import sensor_msgs.point_cloud2 as plc2 
import tf
import numpy as np 

from sensor_msgs.msg import Image, PointCloud2, Imu, NavSatFix
from cv_bridge import CvBridge 
from std_msgs.msg import Header
from visualization_msgs.msg import Marker, MarkerArray
from geometry_msgs.msg import Point, Pose

FRAME_ID = 'map'
LIFETIME = 0.1
DETECTION_CLOLR_DICT = {'Car':(255,255,0), 'Truck':(0,226,255), 'Misc':(141, 40, 255)}

# 矩阵框连接，0，1，2，3是底面的4个点；4，5，6，7顶面的4个点；将矩阵前方以十字岔连接
LINES = [[0, 1], [1, 2], [2, 3], [3, 0]]  # Lower face to Z=0
LINES += [[4, 5], [5, 6], [6, 7], [7, 4]]  # Upper face to Z=0
LINES += [[0, 4], [1, 5], [2, 6], [3, 7]]  # Connections between Lower face and Upper face
LINES += [[4, 1], [5, 0]] # front face

def publish_camera(cam_pub, bridge, image, boxes_2d, types):
	for typ, box in zip(types, boxes_2d):
	    top_left = int(box[0]), int(box[1])
	    bottom_right = int(box[2]), int(box[3])

	    # cv2.rectangle(图片位置,左上角坐标,右下角坐标,颜色,线宽)
	    cv2.rectangle(image, top_left, bottom_right, DETECTION_CLOLR_DICT[typ], 2)
	cam_pub.publish(bridge.cv2_to_imgmsg(image, "bgr8"))

def publish_point_cloud(pcl_pub, point_cloud):
	header = Header()
	header.stamp = rospy.Time.now()
	header.frame_id = FRAME_ID
	pcl_pub.publish(plc2.create_cloud_xyz32(header, point_cloud[:, :3]))

def publish_ego_car(ego_pub):
	marker_array = MarkerArray()

	marker1 = Marker()
	marker1.header.frame_id = FRAME_ID
	marker1.header.stamp = rospy.Time.now()

	marker1.id = 0
	marker1.action = Marker.ADD
	marker1.lifetime = rospy.Duration()
	marker1.type = Marker.LINE_STRIP

	# 初始化四元数为单位四元数（没有旋转）
	marker1.pose.orientation.w = 1.0  # w是实部，对于单位四元数来说，w=1, x=y=z=0
	marker1.pose.orientation.x = 0.0
	marker1.pose.orientation.y = 0.0
	marker1.pose.orientation.z = 0.0

	marker1.color.r = 0.0
	marker1.color.g = 1.0
	marker1.color.b = 0.0
	marker1.color.a = 1.0
	marker1.scale.x = 0.2

	marker1.points = []
	marker1.points.append(Point(10, -10, 0))
	marker1.points.append(Point(0, 0, 0))
	marker1.points.append(Point(10, 10, 0))

	marker_array.markers.append(marker1)

	marker2 = Marker()
	marker2.header.frame_id = FRAME_ID
	marker2.header.stamp - rospy.Time.now()

	marker2.id = -1
	marker2.lifetime = rospy.Duration()
	marker2.type = Marker.MESH_RESOURCE
	marker2.mesh_resource = "package://kitti_tutorial/bmw_x5/BMW X5 4.dae"

	marker2.pose.position.x = 0.0
	marker2.pose.position.y = 0.0
	marker2.pose.position.z = -1.73

	q = tf.transformations.quaternion_from_euler(np.pi/2, 0, np.pi)
	marker2.pose.orientation.x = q[0]
	marker2.pose.orientation.y = q[1]
	marker2.pose.orientation.z = q[2]
	marker2.pose.orientation.w = q[3]

	marker2.color.r = 1.0
	marker2.color.g = 1.0
	marker2.color.b = 1.0
	marker2.color.a = 1.0

	marker2.scale.x = 0.9
	marker2.scale.y = 0.9
	marker2.scale.z = 0.9

	marker_array.markers.append(marker2)

	ego_pub.publish(marker_array)

def publish_imu(imu_pub, imu_data):
	imu = Imu()
	imu.header.frame_id = FRAME_ID
	imu.header.stamp = rospy.Time.now()

	#偏移量
	roll, pitch, yaw = float(imu_data.roll.iloc[0]), float(imu_data.pitch.iloc[0]), float(imu_data.yaw.iloc[0])
	q = tf.transformations.quaternion_from_euler(roll, pitch, yaw)
	imu.orientation.x = q[0]
	imu.orientation.y = q[1]
	imu.orientation.z = q[2]
	imu.orientation.w = q[3]

	#线性加速度
	af, al, au = float(imu_data.af.iloc[0]), float(imu_data.al.iloc[0]), float(imu_data.au.iloc[0])
	imu.linear_acceleration.x = af
	imu.linear_acceleration.y = al
	imu.linear_acceleration.z = au

	# 角速度
	wf, wl, wu = float(imu_data.wf.iloc[0]), float(imu_data.wl.iloc[0]), float(imu_data.wu.iloc[0])
	imu.angular_velocity.x = wf
	imu.angular_velocity.y = wl
	imu.angular_velocity.z = wu

	imu_pub.publish(imu)

def publish_gps(gps_pub, imu_data):
	gps = NavSatFix()
	gps.header.frame_id = FRAME_ID
	gps.header.stamp = rospy.Time.now()

	gps.latitude = imu_data.lat
	gps.longitude = imu_data.lon
	gps.altitude = imu_data.alt

	gps_pub.publish(gps)

def publish_3dbox(box3d_pub, corners_3d_velos, types):
	marker_array = MarkerArray()
	for i, corners_3d_velo in enumerate(corners_3d_velos):
		"""3d检测框"""
		marker = Marker()
		marker.header.frame_id = FRAME_ID
		marker.header.stamp = rospy.Time.now()

		marker.id = i
		marker.action = Marker.ADD
		marker.lifetime = rospy.Duration(LIFETIME)
		marker.type = Marker.LINE_LIST

		# 初始化pose
		marker.pose = Pose()
		marker.pose.position.x = 0.0  # 你可以根据需要设置x, y, z坐标
		marker.pose.position.y = 0.0
		marker.pose.position.z = 0.0

		# 初始化四元数为单位四元数（没有旋转）
		marker.pose.orientation.w = 1.0  # w是实部，对于单位四元数来说，w=1, x=y=z=0
		marker.pose.orientation.x = 0.0
		marker.pose.orientation.y = 0.0
		marker.pose.orientation.z = 0.0

		# 检测框颜色定义
		b, g, r = DETECTION_CLOLR_DICT[types[i]]
		marker.color.r = r/255.0
		marker.color.g = g/255.0
		marker.color.b = b/255.0

		marker.color.a = 1.0
		marker.scale.x = 0.1

		marker.points = []
		for l in LINES:
			p1 = corners_3d_velo[l[0]]
			marker.points.append(Point(p1[0], p1[1], p1[2]))
			p2 = corners_3d_velo[l[1]]
			marker.points.append(Point(p2[0], p2[1], p2[2]))
		marker_array.markers.append(marker)

	box3d_pub.publish(marker_array)

def publish_track(track_pub, tracker, centers):
	marker_array = MarkerArray()

	for track_id in centers:
		marker = Marker()
		marker.header.frame_id = FRAME_ID
		marker.header.stamp = rospy.Time.now()

		marker.action = Marker.ADD
		marker.lifetime = rospy.Duration(LIFETIME)
		marker.type = Marker.LINE_STRIP
		marker.id = track_id

		# 初始化pose
		marker.pose = Pose()
		marker.pose.position.x = 0.0  # 你可以根据需要设置x, y, z坐标
		marker.pose.position.y = 0.0
		marker.pose.position.z = 0.0

		# 初始化四元数为单位四元数（没有旋转）
		marker.pose.orientation.w = 1.0  # w是实部，对于单位四元数来说，w=1, x=y=z=0
		marker.pose.orientation.x = 0.0
		marker.pose.orientation.y = 0.0
		marker.pose.orientation.z = 0.0

		marker.color.r = 1.0
		marker.color.g = 1.0
		marker.color.b = 0.0
		marker.color.a = 0.8
		marker.scale.x = 0.1

		marker.points = []
		for p in tracker[track_id].locations:
			marker.points.append(Point(p[0], p[1], 0))

		marker_array.markers.append(marker)
	track_pub.publish(marker_array)

def publish_distance(distance_pub, minPQDs):
	marker_array = MarkerArray()

	for i, (minP, minQ, minD) in enumerate(minPQDs):
		marker = Marker()
		marker.header.frame_id = FRAME_ID
		marker.header.stamp = rospy.Time.now()

		marker.action = Marker.ADD
		marker.lifetime = rospy.Duration(LIFETIME)
		marker.type = Marker.LINE_STRIP
		marker.id = i 

		# 初始化pose
		marker.pose = Pose()
		marker.pose.position.x = 0.0  # 你可以根据需要设置x, y, z坐标
		marker.pose.position.y = 0.0
		marker.pose.position.z = 0.0

		# 初始化四元数为单位四元数（没有旋转）
		marker.pose.orientation.w = 1.0  # w是实部，对于单位四元数来说，w=1, x=y=z=0
		marker.pose.orientation.x = 0.0
		marker.pose.orientation.y = 0.0
		marker.pose.orientation.z = 0.0

		marker.color.r = 1.0
		marker.color.g = 0.0
		marker.color.b = 1.0
		marker.color.a = 0.5
		marker.scale.x = 0.1

		marker.points = []
		marker.points.append(Point(minP[0], minP[1], 0))
		marker.points.append(Point(minQ[0], minQ[1], 0))

		marker_array.markers.append(marker)

		test_marker = Marker()
		test_marker.header.frame_id = FRAME_ID
		test_marker.header.stamp = rospy.Time.now()

		test_marker.action = Marker.ADD
		test_marker.lifetime = rospy.Duration(LIFETIME)
		test_marker.type = Marker.TEXT_VIEW_FACING
		test_marker.id = i + 1000

		p = (minP + minQ) / 2.0
		test_marker.pose.position.x = p[0]
		test_marker.pose.position.y = p[1]
		test_marker.pose.position.z = 0.0

		test_marker.text = '%.2f'%minD

		test_marker.scale.x = 1
		test_marker.scale.y = 1
		test_marker.scale.z = 1

		test_marker.color.r = 1.0
		test_marker.color.g = 1.0
		test_marker.color.b = 1.0
		test_marker.color.a = 0.8

		marker_array.markers.append(test_marker)

	distance_pub.publish(marker_array)

